import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import csv

from util.util_image import draw_img_xyz
from util.util_csv import read_csv_to_numpy_array
from util.util_ris_pattern_2 import phase_2_pattern_xyz


# 3d画图
def plot_3d_scatter(x_coords, y_coords, z_matrix):
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')

    X, Y = np.meshgrid(x_coords, y_coords)
    Z = z_matrix.T

    ax.scatter(X, Y, Z, c=Z, cmap='gist_rainbow')
    # ax.set_xlabel('X Coordinate')
    # ax.set_ylabel('Y Coordinate')
    # ax.set_zlabel('Z Value')
    # plt.title('3D Scatter Plot')
    # 隐藏轴标签和刻度
    # ax.set_xticks([])
    # ax.set_yticks([])
    # ax.set_zticks([])
    # 隐藏轴标签
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_zlabel('')
    # 显示网格线
    ax.grid(True)
    plt.show()


def draw_img_xyz_3d(data, x, y):
    """
    将2D数据绘制为3D表面图。

    参数:
    data : 2D numpy array
        形成表面高度 (Z值) 和颜色的数据。
    x : 1D or 2D numpy array
        X坐标。如果1D，其长度必须与data的列数匹配。
    y : 1D or 2D numpy array
        Y坐标。如果1D，其长度必须与data的行数匹配。
    """
    # 确保x, y是网格形式，如果不是，则创建网格
    if x.ndim == 1 and y.ndim == 1:
        X, Y = np.meshgrid(x, y)
    elif x.ndim == 2 and y.ndim == 2:
        X, Y = x, y
    else:
        raise ValueError("x和y必须是1D或2D数组，并且形状要与data兼容。")

    Z = data

    # 创建3D绘图对象
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')

    # 绘制3D表面图
    # cmap定义颜色映射，可以根据喜好更改，例如 'viridis', 'plasma', 'jet', 'hot'
    surf = ax.plot_surface(X, Y, Z, cmap='plasma', edgecolor='none', antialiased=True)

    # 添加颜色条
    fig.colorbar(surf, ax=ax, shrink=0.5, aspect=5, label='Pattern Magnitude')

    # 设置图表标题和轴标签
    # ax.set_title('3D Surface Plot')
    # ax.set_xlabel('X-axis')
    # ax.set_ylabel('Y-axis')
    # ax.set_zlabel('Z-axis (Magnitude)')
    # 隐藏轴标签和刻度
    # ax.set_xticks([])
    # ax.set_yticks([])
    # ax.set_zticks([])

    # 改善视角（可选）
    ax.view_init(elev=30, azim=-60) # elev: 仰角, azim: 方位角

    # 显示图形
    plt.tight_layout()
    plt.show()
    # 如果需要保存图片，取消下面一行的注释
    # plt.savefig('3d_surface_plot.png', dpi=200, bbox_inches='tight')


def plot_pattern3d_python(path_csv_phaseBit):
    # 读取码阵
    phaseBit = read_csv_to_numpy_array(path_csv_phaseBit)
    # 计算方向图
    pattern, x, y, z = phase_2_pattern_xyz(phaseBit)
    pattern = np.abs(pattern)
    # 画2d验证
    draw_img_xyz(pattern, x, y)
    # 画3d方向图
    draw_img_xyz_3d(pattern, x, y)
    # plot_3d_scatter(x, y, pattern)


if __name__ == "__main__":
    plot_pattern3d_python(
        "../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/theta(0,60)-phi(0,0)-20250816/phaseBit_AGA_(0,0).csv")
    # plot_pattern3d_python(
    #     "../files/dissertation/chapter_2-sim-AGA/32x32-d9.3mm-FD0.905/beam2-(20,45)-(20,225)/phaseBit_AGA_(((20,45),(20,225))).csv")




